import os
import uuid

import cv2
import numpy as np
import tensorflow as tf


class IdcardRange:
    interpreter = None
    isDebuger = False

    def __init__(self, modelPath, isDebuger):
        # Load the TFLite model and allocate tensors.
        print("loading model ....")
        self.interpreter = tf.lite.Interpreter(model_path=modelPath)
        self.interpreter.allocate_tensors()
        print("loading model .... ok")
        self.isDebuger = isDebuger

    def __detect(self, interpreter, img):
        input_details = interpreter.get_input_details()
        output_details = interpreter.get_output_details()
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        width = int(input_details[0]['shape'][1])
        heigth = int(input_details[0]['shape'][2])
        print("input image  width = {} heigth = {}".format(width, heigth))
        img_rgb = cv2.resize(img_rgb, (width, heigth))
        input_tensor = tf.convert_to_tensor(np.expand_dims(img_rgb, 0), dtype=tf.float32)
        # 这个计算公式
        # ssd_mobilenet_v2_fpn_keras_feature_extractor.py
        # 参考这个文件的 preprocess 方法，计算方式 和 640 的 略有区别
        if width == 320:
            input_tensor = (2.0 / 255.0) * input_tensor - 1.0

        interpreter.set_tensor(input_details[0]['index'], input_tensor)
        interpreter.invoke()
        # 目标区域
        boxes = interpreter.get_tensor(output_details[1]['index'])[0]
        # 评分
        scores = interpreter.get_tensor(output_details[0]['index'])[0]
        classes = interpreter.get_tensor(output_details[3]['index'])[0]
        num_det = interpreter.get_tensor(output_details[2]['index'])[0]
        return boxes, classes, scores, img_rgb

    def __getRect(self, boxes, classes, scores, tmpImg, min_score_thresh=0.6):
        if boxes is not None and len(boxes) <= 0:
            return None, None
        box = None
        for index in range(len(scores)):
            if scores[index] >= min_score_thresh:
                box = tuple(boxes[index].tolist())
                print("box = {}".format(box))
                break
        """
         函数是如何计算出的 目标区域
         viz_utils.visualize_boxes_and_labels_on_image_array
         (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                          ymin * im_height, ymax * im_height)
        """
        if box is None:
            return None, None

        im_height = len(tmpImg)
        im_width = len(tmpImg[0])
        ymin, xmin, ymax, xmax = box
        # 计算坐标
        (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                      ymin * im_height, ymax * im_height)
        # 如果不转 int 会报错
        x = int(left)
        y = int(top)
        w = int(abs(right - left))
        h = int(abs(bottom - top))
        tmpList = []
        xyList = [x, y, w, h]
        # 顺序 左上角  右上角  右下角  左下角
        rect = np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]])
        tmpList.append(rect)
        return tmpList, xyList

    # 预测图片
    def predictByImg(self, tmpImg,min_score_thresh):
        tmpPath = "/tmp/tmp_" + str(uuid.uuid1()) + ".jpg"
        tmpImg.save(tmpPath)
        tmpImg = cv2.imdecode(np.fromfile(tmpPath, dtype=np.uint8), cv2.IMREAD_COLOR)
        #tmpImg = cv2.GaussianBlur(tmpImg, (11, 11), 0)
        resultImg = tmpImg.copy()
        showImg = tmpImg.copy()
        os.remove(tmpPath)
        boxes, classes, scores, img = self.__detect(self.interpreter, tmpImg)
        rectList, xyList = self.__getRect(boxes, classes, scores, tmpImg,min_score_thresh=min_score_thresh)

        # 注意 此处的 rectList 计算的是 原始的 range
        if rectList is not None and len(rectList) > 0:
            #  xyList = [x, y, w, h]
            x = xyList[0]
            y = xyList[1]
            w = xyList[2]
            h = xyList[3]
            resultImg = resultImg[y:y + h, x:x + w]
        else:
            resultImg = None

        tmp = cv2.drawContours(showImg, rectList, -1, color=(0, 0, 255), thickness=3)
        imgName = str(uuid.uuid1()) + ".jpg"
        cv2.imwrite("/tmp/" + imgName, tmp)

        if rectList is not None and self.isDebuger:
            cv2.imshow("frame", tmpImg)
            cv2.waitKey(0)
            cv2.destroyWindow("frame")
        return boxes, classes, scores, xyList, resultImg, imgName
